import cv2
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import gridspec
from matplotlib.font_manager import FontProperties


def phase_correlation(a, b):
    """计算相位相关"""
    ffta = np.fft.fft2(a)
    fftb = np.fft.fft2(b)
    product = ffta * np.conj(fftb)
    quotient = product / (np.abs(product) + 1e-8)  # 避免除零
    cc = np.fft.ifft2(quotient)
    return np.abs(np.fft.fftshift(cc))


def fingerprint_registration(img1, img2):
    # 细节点检测（SIFT）
    sift = cv2.SIFT_create()
    kp1, des1 = sift.detectAndCompute(img1, None)
    kp2, des2 = sift.detectAndCompute(img2, None)

    # FLANN匹配
    FLANN_INDEX_KDTREE = 1
    index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
    search_params = dict(checks=50)
    flann = cv2.FlannBasedMatcher(index_params, search_params)
    matches = flann.knnMatch(des1, des2, k=2)

    # 筛选匹配点
    good_matches = []
    for m, n in matches:
        if m.distance < 0.7 * n.distance:
            good_matches.append(m)

    if len(good_matches) < 4:
        raise ValueError("匹配点不足")

    # 粗配准（单应性变换）
    src_pts = np.float32([kp1[m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2)
    dst_pts = np.float32([kp2[m.trainIdx].pt for m in good_matches]).reshape(-1, 1, 2)
    M, _ = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)
    warped_img = cv2.warpPerspective(img1, M, (img2.shape[1], img2.shape[0]))

    # 精细配准（相位相关）
    h, w = img2.shape
    warped_roi = warped_img[:h, :w]
    cc = phase_correlation(warped_roi, img2)
    _, _, _, max_loc = cv2.minMaxLoc(cc)
    dx, dy = max_loc[0] - w // 2, max_loc[1] - h // 2

    M_shift = np.float32([[1, 0, dx], [0, 1, dy]])
    fine_warped = cv2.warpAffine(warped_img, M_shift, (w, h))

    # 计算匹配分数
    score = np.sum(fine_warped * img2) / np.sqrt(np.sum(fine_warped ** 2) * np.sum(img2 ** 2))
    # 绘制特征匹配图
    match_img = cv2.drawMatches(img1, kp1, img2, kp2, good_matches, None,
                                flags=cv2.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS)
    return warped_img, fine_warped, score, match_img


def fingerprint_inpainting(img):
    # 图像预处理，转换为二值图像
    _, binary = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY_INV)
    # 形态学操作，填充小的空洞
    kernel = np.ones((3, 3), np.uint8)
    opening = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel, iterations=2)
    # 查找轮廓
    contours, _ = cv2.findContours(opening.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    mask = np.zeros_like(img)
    # 绘制轮廓到掩膜上
    cv2.drawContours(mask, contours, -1, 255, -1)
    # 使用图像修复技术修复指纹
    inpainted = cv2.inpaint(img, mask, 3, cv2.INPAINT_TELEA)
    return inpainted


# 读取图像（需替换为实际路径）
img1 = cv2.imread('f3.jpg', 0)
img2 = cv2.imread('f2.jpg', 0)

# 检查图像是否成功读取
if img1 is None or img2 is None:
    print("图像读取失败，请检查文件路径和文件完整性。")
else:
    try:
        warped_img, fine_warped, match_score, match_img = fingerprint_registration(img1, img2)
        print(f"匹配分数: {match_score}")

        # 加载本地中文字体
        font_path = 'fonts/AlibabaPuHuiTi-Regular.ttf'  # 替换为你的中文字体文件路径
        font = FontProperties(fname=font_path)

        # 对精细配准后的指纹进行修复
        repaired_fingerprint = fingerprint_inpainting(fine_warped)

        # 设定判别阈值
        threshold = 0.8  # 可根据实际情况调整
        if match_score >= threshold:
            result = "是同一个人的同一根手指"
        else:
            result = "不是同一个人的同一根手指"
        print(f"判别结果: {result}")

        # 设置图片显示
        fig = plt.figure(figsize=(18, 12))
        gs = gridspec.GridSpec(2, 3, width_ratios=[1, 1, 1], height_ratios=[1, 1])

        # 显示原始指纹1
        ax0 = plt.subplot(gs[0, 0])
        ax0.imshow(img1, cmap='gray')
        ax0.set_title('原始指纹1', fontproperties=font)
        ax0.axis('off')

        # 显示原始指纹2
        ax1 = plt.subplot(gs[0, 1])
        ax1.imshow(img2, cmap='gray')
        ax1.set_title('原始指纹2', fontproperties=font)
        ax1.axis('off')

        # 显示特征匹配图
        ax2 = plt.subplot(gs[0, 2])
        ax2.imshow(cv2.cvtColor(match_img, cv2.COLOR_BGR2RGB))
        ax2.set_title('特征匹配图', fontproperties=font)
        ax2.axis('off')

        # 显示粗配准指纹
        ax3 = plt.subplot(gs[1, 0])
        ax3.imshow(warped_img, cmap='gray')
        ax3.set_title('粗配准指纹', fontproperties=font)
        ax3.axis('off')

        # 显示精细配准指纹
        ax4 = plt.subplot(gs[1, 1])
        ax4.imshow(fine_warped, cmap='gray')
        ax4.set_title('精细配准指纹', fontproperties=font)
        ax4.axis('off')

        # 显示最终处理修正的完整指纹图像
        ax5 = plt.subplot(gs[1, 2])
        ax5.imshow(repaired_fingerprint, cmap='gray')
        ax5.set_title('最终处理修正指纹', fontproperties=font)
        ax5.axis('off')

        plt.tight_layout()
        plt.show()

    except ValueError as e:
        print(e)
